package de.distMLP.communication;

import java.io.IOException;
import java.util.Map;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hama.bsp.BSPPeer;
import org.apache.hama.bsp.sync.SyncException;

import de.distMLP.messages.HamaMessage;
import de.distMLP.messages.WeightMatrixMessage;
import de.mlp_distributed.mlp.math.mahout.Matrix;

public class BtreeMessageDistributor extends MessageDistributor {

	private static final Log LOG = LogFactory.getLog(BtreeMessageDistributor.class);

	public BtreeMessageDistributor(final BSPPeer<?, ?, ?, ?, HamaMessage> peer, final String masterPeer) {
		super(peer, masterPeer);
	}

	@Override
	public HamaMessage distributeMessageToAllPeers(final String senderName, final HamaMessage msg, final String[] peerNames)
			throws IOException, SyncException, InterruptedException {
		HamaMessage tmp = msg;

		final BSPPeer<?, ?, ?, ?, HamaMessage> peer = this.getPeer();

		Map<Integer, String> peerMap;
		Map<String, Integer> peerToIntMap;

		if (this.isLocalMode()) {
			peerMap = this.getLocalPeerMap();
			peerToIntMap = this.getLocalPeerToIntMap();
		} else {
			peerMap = this.getPeerMap();
			peerToIntMap = this.getPeerToIntMap();
		}

		if (peer.getPeerName().equals(this.getMasterPeer())) {
			final String peerNameIteration1 = BtreeMessageDistributor.getDataReceiver(peerMap, 1, 1);
			if (peerNameIteration1 != null) {
				peer.send(peerNameIteration1, msg);
			}
		}
		peer.sync();
		final int nbComIterations = BtreeMessageDistributor.calculateNbOfIterations(peerMap.keySet().size());

		BtreeMessageDistributor.LOG.info("One-To-All: " + peer.getNumCurrentMessages() + " messages on peer " + peer.getPeerName() + "("
				+ peerToIntMap.get(peer.getPeerName()) + ") Iteration:" + nbComIterations);
		/**
		 * At this point the master already sent the weight matrix to one peer.
		 */

		for (int i = 2; i <= nbComIterations; i++) {
			if ((tmp == null) && (peer.getNumCurrentMessages() > 0)) {
				tmp = peer.getCurrentMessage();
			}
			if (tmp != null) {
				final String peerName = BtreeMessageDistributor.getDataReceiver(peerMap, peerToIntMap.get(peer.getPeerName()), i);
				if (peerName != null) {
					peer.send(peerName, tmp);
				}
			}
			peer.sync();

			BtreeMessageDistributor.LOG.info("One-To-All: " + peer.getNumCurrentMessages() + " messages on peer " + peer.getPeerName()
					+ "(" + peerToIntMap.get(peer.getPeerName()) + ")" + "Iteration: " + i);
		}
		if (!this.isLocalMode()) {
			/**
			 * At this point the first thread on each server is having the
			 * weight matrix message.
			 */
			BtreeMessageDistributor.LOG.info("Local-One-To-All: " + peer.getNumCurrentMessages() + " messages on peer "
					+ peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")");
			if ((tmp == null) && (peer.getNumCurrentMessages() > 0)) {
				tmp = peer.getCurrentMessage();
			}

			if (tmp != null) {
				for (final String peerName : this.getLocalPeerToIntMap().keySet()) {
					if (peerName.equals(peer.getPeerName())) {
						continue;
					}
					peer.send(peerName, tmp);
				}
			}
			peer.sync();
			BtreeMessageDistributor.LOG.info("Local-One-To-All: " + peer.getNumCurrentMessages() + " messages on peer "
					+ peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")");
			if ((tmp == null) && (peer.getNumCurrentMessages() > 0)) {
				tmp = peer.getCurrentMessage();
			}
		}
		return tmp;
	}

	@Override
	public HamaMessage distributeAndAggegateMessageToMasterPeer(final HamaMessage msg, final String masterPeer) throws IOException,
			SyncException, InterruptedException {
		final BSPPeer<?, ?, ?, ?, HamaMessage> peer = this.getPeer();

		if (peer.getNumPeers() == 1) {
			return msg;
		}

		WeightMatrixMessage averageSumMsg = null;

		final WeightMatrixMessage gradientSumMsg = (WeightMatrixMessage) msg;

		Map<Integer, String> peerMap;
		Map<String, Integer> peerToIntMap;

		if (this.isLocalMode()) {
			peerMap = this.getLocalPeerMap();
			peerToIntMap = this.getLocalPeerToIntMap();
		} else {
			peerMap = this.getPeerMap();
			peerToIntMap = this.getPeerToIntMap();
		}

		if (!this.isLocalMode()) {
			this.sendMsgToFirstPeerOnLocalhost(msg, peerToIntMap);
			peer.sync();
		}

		final int nbComIterations = BtreeMessageDistributor.calculateNbOfIterations(peerMap.keySet().size());

		// Der erste Thread hat jetzt x-Nachrichten
		if (peer.getNumCurrentMessages() > 0) {
			BtreeMessageDistributor.LOG.info("Agg. Messages: " + peer.getNumCurrentMessages() + " messages on peer " + peer.getPeerName()
					+ "(" + peerToIntMap.get(peer.getPeerName()) + ")");
			averageSumMsg = this.calculateWeightSum(peer, gradientSumMsg);
			final String peerName = BtreeMessageDistributor.getReverseDataReceiver(peerMap, peerToIntMap.get(peer.getPeerName()),
					nbComIterations);
			if (peerName != null) {
				peer.send(peerName, averageSumMsg);

				BtreeMessageDistributor.LOG.info(peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")"
						+ " sends average weight matrix to " + peerName + ". Iteration: " + nbComIterations);
			} else {
				BtreeMessageDistributor.LOG.info(peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")"
						+ " Got no incoming message and no peer to send to. Iteration: " + nbComIterations);
			}
		} else {
			if ((this.getNbLocalThreads(peer) == 1) || this.isLocalMode()) {
				// In case of only one thread per peer there are no other
				// messages
				averageSumMsg = gradientSumMsg;
				final String peerName = BtreeMessageDistributor.getReverseDataReceiver(peerMap, peerToIntMap.get(peer.getPeerName()),
						nbComIterations);
				if ((peerName != null) && (gradientSumMsg != null)) {
					peer.send(peerName, averageSumMsg);
				}
				BtreeMessageDistributor.LOG.info("no messages on peer " + peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName())
						+ ")" + ". Sending message to peer " + peerName);
			} else {
				BtreeMessageDistributor.LOG.info("no messages on peer " + peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName())
						+ ")");
			}
		}
		peer.sync();

		for (int i = nbComIterations - 1; i >= 1; i--) {

			String peerName = null;
			if (peerToIntMap.containsKey(peer.getPeerName())) {
				peerName = BtreeMessageDistributor.getReverseDataReceiver(peerMap, peerToIntMap.get(peer.getPeerName()), i);
			}

			if ((averageSumMsg != null) && (peerName != null)) {
				if (peer.getNumCurrentMessages() > 0) {
					averageSumMsg = this.calculateWeightSum(peer, averageSumMsg);
					BtreeMessageDistributor.LOG.info(peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")"
							+ " sends average weight matrix to " + peerName + ". Iteration: " + i);

					peer.send(peerName, averageSumMsg);
				} else {
					peer.send(peerName, averageSumMsg);

					BtreeMessageDistributor.LOG.info(peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")"
							+ " got no incomming messages but a peer to send to. " + ". Iteration: " + i);
				}
			} else {
				if (peer.getNumCurrentMessages() > 0) {
					averageSumMsg = this.calculateWeightSum(peer, averageSumMsg);
				} else {
					BtreeMessageDistributor.LOG.info(peer.getPeerName() + "(" + peerToIntMap.get(peer.getPeerName()) + ")"
							+ " Got no incoming message and no peer to send to. Iteration: " + i);
				}
			}
			peer.sync();
		}
		if (peer.getPeerName().equals(masterPeer)) {
			averageSumMsg = this.calculateWeightSum(peer, averageSumMsg);
		}
		return averageSumMsg;
	}

	private void sendMsgToFirstPeerOnLocalhost(final HamaMessage msg, final Map<String, Integer> peerToIntMap) throws IOException {
		final BSPPeer<?, ?, ?, ?, HamaMessage> peer = this.getPeer();

		final String hostName = peer.getPeerName().split(":")[0];

		for (final String peerName : peerToIntMap.keySet()) {
			// TODO code zu bestimmung des ersten threads in setup
			// auslagern
			if (peerName.equals(peer.getPeerName())) {
				// Don't send the message to yourselves
				break;
			}
			final String tmpHostname = peerName.split(":")[0];
			if (tmpHostname.equals(hostName)) {
				peer.send(peerName, msg);
				break;
			}
		}
	}

	private int getNbLocalThreads(final BSPPeer<?, ?, ?, ?, HamaMessage> peer) {
		int counter = 0;
		final String hostName = peer.getPeerName().split(":")[0];

		for (final String peerName : peer.getAllPeerNames()) {
			final String tmpHostname = peerName.split(":")[0];
			if (tmpHostname.equals(hostName)) {
				counter++;
			}
		}
		return counter;
	}

	protected WeightMatrixMessage calculateWeightSum(final BSPPeer<?, ?, ?, ?, HamaMessage> peer, final WeightMatrixMessage other)
			throws IOException {
		Matrix[] weightChangesSum = null;

		int totalNumberOfMatrices = peer.getNumCurrentMessages();

		if (totalNumberOfMatrices == 0) {
			BtreeMessageDistributor.LOG.error("Iteration failed. No messages to process on master.");
			return null;
		}
		boolean stopTraining = false;

		// total sum of errors
		double cost = 0;

		if (other != null) {
			if (other.isStopTraining()) {
				stopTraining = true;
			}
			final Matrix[] matrix = other.getMatrix();
			if (matrix == null) {
				totalNumberOfMatrices--;
			} else {
				if (weightChangesSum == null) {
					weightChangesSum = new Matrix[matrix.length];
				}
				totalNumberOfMatrices++;
				// sum up all weight changes
				for (int i = 0; i < matrix.length; i++) {
					if (weightChangesSum[i] == null) {
						weightChangesSum[i] = matrix[i].clone();
					} else {
						weightChangesSum[i] = weightChangesSum[i].plus(matrix[i]);
					}
				}

				cost += other.getCost();

			}
		}

		while (peer.getNumCurrentMessages() > 0) {
			try {
				final WeightMatrixMessage msg = (WeightMatrixMessage) peer.getCurrentMessage();

				if (msg.isStopTraining()) {
					stopTraining = true;
				}
				final Matrix[] matrix = msg.getMatrix();
				if (matrix == null) {
					totalNumberOfMatrices--;
					continue;
				}
				if (weightChangesSum == null) {
					weightChangesSum = new Matrix[matrix.length];
				}
				// sum up all weight changes
				for (int i = 0; i < matrix.length; i++) {
					if (weightChangesSum[i] == null) {
						weightChangesSum[i] = matrix[i].clone();
					} else {
						weightChangesSum[i] = weightChangesSum[i].plus(matrix[i]);
					}
				}
				cost += msg.getCost();

			} catch (final IOException e) {
				BtreeMessageDistributor.LOG.error(e, e);
			}
		}
		return new WeightMatrixMessage(weightChangesSum, stopTraining, cost);
	}

	/**
	 * @param totalNbPeers
	 *            total number of peers
	 * @param peerMap
	 *            map of peer names
	 * @param peerName
	 *            number of this peer (starting with one)
	 * @param iterationNumber
	 *            number of iteration (starting with one)
	 * @return name of peer to send data to or null
	 */
	@SuppressWarnings("boxing")
	private static String getDataReceiver(final Map<Integer, String> peerMap, final int peerNumber, final int iterationNumber) {
		final int totalNbPeers = peerMap.keySet().size();
		String receiver = null;
		final int sender = (int) Math.pow(2, iterationNumber - 1);
		if (peerNumber <= sender) {
			final int receiverNumber = peerNumber + sender;
			if (receiverNumber <= totalNbPeers) {
				receiver = peerMap.get(receiverNumber);
			}
		}
		return receiver;
	}

	/**
	 * @param totalNbPeers
	 *            total number of peers
	 * @param peerMap
	 *            map of peer names
	 * @param peerName
	 *            number of this peer (starting with one)
	 * @param iterationNumber
	 *            number of iteration (starting with one)
	 * @return name of peer to send data to or null
	 */
	@SuppressWarnings("boxing")
	private static String getReverseDataReceiver(final Map<Integer, String> peerMap, final int peerNumber, final int iterationNumber) {
		String receiver = null;
		final int sender = (int) Math.pow(2, iterationNumber - 1);
		final int senderMax = (int) Math.pow(2, iterationNumber);
		if ((peerNumber > sender) && (peerNumber <= senderMax)) {
			final int receiverNumber = peerNumber - sender;
			if (receiverNumber > 0) {
				receiver = peerMap.get(receiverNumber);
			}
		}
		return receiver;
	}

	private static int calculateNbOfIterations(final int numberOfPeers) {
		final double log = Math.log(numberOfPeers) / Math.log(2);
		final int tmpLog = (int) log;

		if (tmpLog < log) {
			return tmpLog + 1;
		}
		return tmpLog;
	}

}
